#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
MongoShake 容器内定时同步脚本
简化版本，专门在容器内运行
"""

import os
import sys
import time
import logging
import subprocess
import signal
from datetime import datetime
import json

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('/var/log/mongoshake/container_scheduler.log'),
        logging.StreamHandler(sys.stdout)
    ]
)
logger = logging.getLogger(__name__)

class ContainerScheduler:
    def __init__(self):
        self.running = True
        self.load_config()
        
        # 设置信号处理
        signal.signal(signal.SIGINT, self.signal_handler)
        signal.signal(signal.SIGTERM, self.signal_handler)
    
    def load_config(self):
        """加载配置文件"""
        try:
            with open('/app/scheduler_config.json', 'r', encoding='utf-8') as f:
                self.config = json.load(f)
        except FileNotFoundError:
            # 默认配置
            self.config = {
                "sync_interval_hours": 5,
                "max_retries": 3,
                "retry_delay": 60,
                "sync_timeout_minutes": 30
            }
    
    def signal_handler(self, signum, frame):
        """信号处理函数"""
        logger.info(f"收到信号 {signum}，正在停止调度器...")
        self.running = False
    
    def check_mongoshake_process(self):
        """检查MongoShake进程是否运行"""
        try:
            result = subprocess.run(
                ["pgrep", "-f", "collector"],
                capture_output=True,
                text=True
            )
            return result.returncode == 0
        except Exception as e:
            logger.warning(f"检查MongoShake进程失败: {e}")
            return False
    
    def restart_mongoshake(self):
        """重启MongoShake进程"""
        try:
            logger.info("重启MongoShake进程...")
            
            # 停止当前进程
            subprocess.run(["pkill", "-f", "collector"], timeout=30)
            time.sleep(5)
            
            # 启动新进程
            subprocess.Popen([
                "/app/collector", 
                "-conf", "/etc/mongoshake.conf"
            ])
            
            logger.info("MongoShake进程重启完成")
            return True
            
        except Exception as e:
            logger.error(f"重启MongoShake进程失败: {e}")
            return False
    
    def check_sync_status(self):
        """检查同步状态"""
        try:
            # 检查MongoShake的HTTP状态接口
            result = subprocess.run(
                ["curl", "-s", "http://localhost:9100"],
                capture_output=True,
                text=True,
                timeout=10
            )
            
            if result.returncode == 0:
                logger.info("MongoShake服务状态正常")
                return True
            else:
                logger.warning("MongoShake服务状态异常")
                return False
                
        except Exception as e:
            logger.warning(f"检查同步状态失败: {e}")
            return False
    
    def perform_sync(self):
        """执行同步操作"""
        logger.info("开始执行定时同步...")
        start_time = datetime.now()
        
        try:
            # 检查当前进程状态
            if not self.check_mongoshake_process():
                logger.warning("MongoShake进程未运行，启动进程...")
                subprocess.Popen([
                    "/app/collector", 
                    "-conf", "/etc/mongoshake.conf"
                ])
                time.sleep(10)
            
            # 重启进程以开始新的同步
            if not self.restart_mongoshake():
                logger.error("重启MongoShake进程失败")
                return False
            
            # 等待服务启动
            time.sleep(15)
            
            # 检查同步状态
            retry_count = 0
            while retry_count < self.config["max_retries"]:
                if self.check_sync_status():
                    logger.info("同步状态检查通过")
                    break
                retry_count += 1
                logger.warning(f"同步状态检查失败，重试 {retry_count}/{self.config['max_retries']}")
                time.sleep(self.config["retry_delay"])
            else:
                logger.error("同步状态检查失败，达到最大重试次数")
                return False
            
            # 记录同步完成
            end_time = datetime.now()
            duration = end_time - start_time
            
            logger.info(f"定时同步完成！耗时: {duration}")
            
            return True
            
        except Exception as e:
            logger.error(f"同步过程中出错: {e}")
            return False
    
    def run_once(self):
        """执行一次同步"""
        logger.info("执行单次同步...")
        return self.perform_sync()
    
    def run_continuous(self, interval_hours=5):
        """连续运行调度器"""
        logger.info(f"启动连续调度器，间隔: {interval_hours}小时")
        
        while self.running:
            try:
                # 执行同步
                self.perform_sync()
                
                # 等待下次同步
                wait_seconds = interval_hours * 3600
                logger.info(f"等待 {interval_hours} 小时后进行下次同步...")
                
                # 每10分钟检查一次是否需要停止
                check_interval = 600
                for _ in range(wait_seconds // check_interval):
                    if not self.running:
                        break
                    time.sleep(check_interval)
                
                # 剩余等待时间
                remaining = wait_seconds % check_interval
                if remaining > 0 and self.running:
                    time.sleep(remaining)
                
            except KeyboardInterrupt:
                logger.info("收到中断信号，正在停止...")
                break
            except Exception as e:
                logger.error(f"调度器运行出错: {e}")
                time.sleep(60)  # 出错后等待1分钟再继续
        
        logger.info("连续调度器已停止")

def main():
    """主函数"""
    import argparse
    
    parser = argparse.ArgumentParser(description='MongoShake容器内定时同步脚本')
    parser.add_argument('--interval', type=int, default=5, 
                       help='同步间隔（小时），默认5小时')
    parser.add_argument('--once', action='store_true',
                       help='只执行一次同步')
    
    args = parser.parse_args()
    
    # 创建调度器
    scheduler = ContainerScheduler()
    
    if args.once:
        # 只执行一次
        success = scheduler.run_once()
        sys.exit(0 if success else 1)
    else:
        # 连续运行
        scheduler.run_continuous(args.interval)

if __name__ == "__main__":
    main() 